library(tidyverse)
library(maptools)
library(igraph)
library(parallel)
library(redist)
library(spdep)
library(gridExtra)
library(sf)

# prep the map shapes
fl_map <- readShapePoly("../../data/fl/FL.shp")

## Load nodes
mn <- 89 
nodes <- strsplit(readLines("/n/imai_lab/bfifield/enumeration/largemap_enum_all/map_sub_nodes_89.dat")," ")[[1]]

## First, subset down to right map
fl_sub <- fl_map[which(rownames(fl_map@data) %in% nodes),]

## Then, reorder nodes to be correct order to match - otherwise, won't
## line up properly with original map
fl_sub <- fl_sub[match(nodes, rownames(fl_sub@data)),]

## Convert factor columns to characters, because this particular data is messy
indx <- sapply(fl_sub@data, is.factor)
fl_sub@data[indx] <- lapply(
  fl_sub@data[indx], function(x) as.numeric(as.character(x))
)

fl_sub <- st_as_sf(fl_sub)

# 18 'loops'
loop <- as.numeric(Sys.getenv("SLURM_ARRAY_TASK_ID"))
lines_to_read <- seq(from = 0, to = 45000000, by = 2500000)
lines_to_read[length(lines_to_read)] <- 44082156


  i<-loop+1
  ## Load in solutions
  sols_all <- read_lines("/n/imai_lab/ckenny/enum_all/solutions_89.dat",
                         skip = lines_to_read[i-1],
                         n_max = lines_to_read[i] - lines_to_read[i-1])
  
  ## Create edgelist
  el <- strsplit(readLines("/n/imai_lab/bfifield/enumeration/largemap_enum_all/map_sub_ordered_89.dat"), split = " ")
  el <- apply(do.call(rbind, el), 2, as.numeric)
  
  ## All solutions to matrix
  ndists <- 2
  sols_out_all <- mclapply(1:length(sols_all), function(x){
    sol_split <- as.numeric(strsplit(sols_all[x], split = " ")[[1]])
    el_sub <- el[sol_split,]
    comps_out <- components(graph_from_edgelist(el_sub, directed = FALSE))
    if(comps_out$no < ndists){
      max_num <- max(comps_out$membership)
      comps_out <- c(comps_out$membership, (max_num + 1):ndists)
    }else{
      comps_out <- comps_out$membership
    }
    return(comps_out)
  }, mc.cores = detectCores()-1)
  sols_out_all <- do.call(cbind, sols_out_all)
  
  ## Calculate compacntess
pop <- fl_sub$pop*t(matrix(rep(fl_sub$pop,nrow(fl_sub)),nrow(fl_sub)))
centroids <- st_geometry(st_centroid(x = fl_sub))
distances <- st_distance(centroids, centroids)^2

system.time(
out <- unlist(mclapply(1:ncol(sols_out_all), function(x){
 ind <- sols_out_all[, x] == 1
 return(sum(pop[ind,ind]*distances[ind,ind]) + sum(pop[!ind,!ind]*distances[!ind,!ind]))
}, mc.cores = detectCores() - 1)))


comp_sample <- data.frame(fh = out)
write_csv(comp_sample, path = paste0('/n/imai_lab/ckenny/enum_all/fh_',loop,'.csv'))

